package utils

import (
	"debug/elf"
	"fmt"
)

// getElfSymbol 查找 ELF 文件可被hook的符号的地址
func GetElfSymbol(binPath string) ([]string, error) {

	elfFile, err := elf.Open(binPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open ELF file: %w", err)
	}
	defer elfFile.Close()

	symbols := make([]string, 0)

	// 遍历符号表
	syms, err := elfFile.Symbols()
	if err != nil {
		return nil, fmt.Errorf("failed to read symbols: %w", err)
	}

	for _, sym := range syms {
		// 过滤出全局函数符号 && 非未定义符号（动态链接符号）
		if elf.ST_TYPE(sym.Info) == elf.STT_FUNC && sym.Section != elf.SHN_UNDEF {
			symbols = append(symbols, sym.Name)
		}
	}

	return symbols, nil
}

func GetSoSymbol(binPath string) ([]string, error) {

	elfFile, err := elf.Open(binPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open ELF file: %w", err)
	}
	defer elfFile.Close()

	symbols := make([]string, 0)

	// 遍历符号表
	syms, err := elfFile.DynamicSymbols()
	if err != nil {
		return nil, fmt.Errorf("failed to read symbols: %w", err)
	}

	for _, sym := range syms {
		// 过滤出全局函数符号
		if elf.ST_TYPE(sym.Info) == elf.STT_FUNC && sym.Section != elf.SHN_UNDEF {
			symbols = append(symbols, sym.Name)
		}
	}
	return symbols, nil
}

// 找到elfPath中属于soPath库中的函数符号
func GetElfSoSymbol(elfPath, soPath string) ([]string, error) {

	elfFile, err := elf.Open(elfPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open ELF file: %w", err)
	}
	defer elfFile.Close()

	soFile, err := elf.Open(soPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open shared object file: %w", err)
	}
	defer soFile.Close()

	soSyms, err := soFile.DynamicSymbols()
	if err != nil {
		return nil, fmt.Errorf("failed to read symbols from shared object: %w", err)
	}

	soSymSet := make(map[string]struct{})
	for _, sym := range soSyms {
		if elf.ST_TYPE(sym.Info) == elf.STT_FUNC {
			soSymSet[sym.Name] = struct{}{}
		}
	}

	elfSyms, err := elfFile.DynamicSymbols()
	if err != nil {
		return nil, fmt.Errorf("failed to read symbols from ELF file: %w", err)
	}

	symbols := make([]string, 0)

	for _, sym := range elfSyms {
		if elf.ST_TYPE(sym.Info) == elf.STT_FUNC {
			if _, exists := soSymSet[sym.Name]; exists {
				symbols = append(symbols, sym.Name)
			}
		}
	}

	return symbols, nil
}
